import rclpy
from rclpy.node import Node
from origincar_msg.msg import Sign
from sensor_msgs.msg import Image
import cv2
from cv_bridge import CvBridge
from pyzbar.pyzbar import decode
import numpy as np
from ai_msgs.msg import PerceptionTargets
import time
import psutil

class QrcodeNode(Node):
    def __init__(self, name):
        super().__init__(name)
        self.bridge = CvBridge()
        self.get_logger().info("大家好，我是发布二维码识别的话%s" % name)
        
        # 创建发布者，发布到 /sign_switch 话题
        self.publisher_qrcode = self.create_publisher(Sign, 'sign_switch', 10)
        
        # 创建订阅者，订阅 /image 话题（用于接收图像）
        self.subscription = self.create_subscription(
            Image,
            '/image',
            self.image_callback,
            10
        )
        
        # 创建订阅者，订阅 hobot_dnn_detection 话题（用于接收 YOLO 识别结果）
        self.subscription_yolo = self.create_subscription(
            PerceptionTargets,
            'hobot_dnn_detection',
            self.listener_yolo_callback,
            10
        )

        # 初始化变量存储最新的二维码 ROI 和图像
        self.latest_roi = None
        self.latest_image = None

        # 初始化微信二维码检测器
        self.wechat_qr_detector = QRCodeDetector()

    def listener_yolo_callback(self, msg):
        # 遍历检测结果，寻找 QRCode 类型的目标
        for target in msg.targets:
            class_type = target.rois[0].type.strip()
            if class_type == 'qrcode':
                roi = target.rois[0].rect
                # 更新二维码的 ROI 参数
                self.latest_roi = (roi.x_offset, roi.y_offset, roi.width, roi.height)
                self.get_logger().info("检测到二维码区域，准备解码")
            else:
                self.latest_roi = None
                self.get_logger().info("未检测到二维码区域")

    def image_callback(self, msg):
        if self.latest_roi is None:
            return  # 如果没有检测到二维码 ROI，跳过图像处理

        # 将 ROS 图像消息转换为 OpenCV 图像格式
        np_arr = np.frombuffer(msg.data, np.uint8)
        self.latest_image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)

        # 从 ROI 区域裁剪图像，仅对该区域进行解码
        x, y, w, h = self.latest_roi
        if x > 10 and x+w+10 < 640 and y > 10 and y+h+10 < 480:
            roi_image = self.latest_image[y-10:y+h+10, x-10:x+w+10]
        else:
            roi_image = self.latest_image[y:y+h, x:x+w]
        # 注释掉不需要的解码方法
        # qr_codes_pyzbar = self.decode_qr_code_pyzbar(roi_image)  # 使用 pyzbar 解码
        qr_codes_wechat = self.decode_qr_code_wechat(roi_image)  # 使用 WeChatQRCode 解码

        # 处理解码结果
        # self.process_decoded_qr(qr_codes_pyzbar)
        self.process_decoded_qr(qr_codes_wechat)

    def decode_qr_code_pyzbar(self, image):

        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        decoded_objects = decode(gray)
        return decoded_objects

    def decode_qr_code_wechat(self, image):
        
        qr_code_data = self.wechat_qr_detector.detect(image)
        return qr_code_data

    def process_decoded_qr(self, qr_codes):
        if isinstance(qr_codes, list) and qr_codes:
            for obj in qr_codes:
                qr_code_data = obj.data.decode('utf-8')
                if qr_code_data in ["ClockWise", "AntiClockWise"]:
                    self.publish_sign_message(qr_code_data)
                    return  # 立即返回，避免重复处理

        elif isinstance(qr_codes, str) and qr_codes in ["ClockWise", "AntiClockWise"]:
            self.publish_sign_message(qr_codes)

    def publish_sign_message(self, qr_code_data):
        self.get_logger().info(f"检测到二维码: {qr_code_data}")
        sign_msg = Sign()
        sign_msg.sign_data = 3 if qr_code_data == "ClockWise" else 4 if qr_code_data == "AntiClockWise" else 0
        self.publisher_qrcode.publish(sign_msg)

class QRCodeDetector:
    def __init__(self):
        self.depro = '/root/lyz_qrcode/detect.prototxt'
        self.decaf = '/root/lyz_qrcode/detect.caffemodel'
        self.srpro = '/root/lyz_qrcode/sr.prototxt'
        self.srcaf = '/root/lyz_qrcode/sr.caffemodel'

    def detect(self, img):
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        detector = cv2.wechat_qrcode_WeChatQRCode(self.depro, self.decaf, self.srpro, self.srcaf)
        barcodes, _ = detector.detectAndDecode(gray)
        return barcodes[0] if barcodes else None

def main(args=None):
    rclpy.init(args=args)
    node = QrcodeNode("qrcode")  # 创建节点
    rclpy.spin(node)  # 保持节点运行
    rclpy.shutdown()  # 关闭节点

if __name__ == '__main__':
    main()
